import pathlib

import matplotlib.pyplot
import torch

import models
import shared.checkpoint


def main():
    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    checkpoints = checkpoint_manager.existed()
    checkpoints_sorted = sorted(checkpoints.items(), key=lambda x: x[0])

    matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    for i, path in checkpoints_sorted:
        model = models.PinnOfKortewegDeVriesEquation()
        checkpoint = torch.load(path)
        model.load_state_dict(checkpoint["model"])

        matplotlib.pyplot.subplot(221)
        matplotlib.pyplot.ylim(top=1.1)
        matplotlib.pyplot.scatter(i, torch.detach(model.lambda1), s=3, c="tab:blue")
        matplotlib.pyplot.subplot(223)
        matplotlib.pyplot.scatter(i, torch.detach(model.lambda2), s=3, c="tab:orange")

    matplotlib.pyplot.subplot(221)
    matplotlib.pyplot.axhline(1, c="tab:green")
    matplotlib.pyplot.xlabel(r"Epoch")
    matplotlib.pyplot.ylabel(r"$\lambda_1$ Estimated")

    matplotlib.pyplot.subplot(223)
    matplotlib.pyplot.axhline(0.0025, c="tab:green")
    matplotlib.pyplot.xlabel(r"Epoch")
    matplotlib.pyplot.ylabel(r"$\lambda_2$ Estimated")

    for i, path in checkpoints_sorted[-100:]:
        model = models.PinnOfKortewegDeVriesEquation()
        checkpoint = torch.load(path)
        model.load_state_dict(checkpoint["model"])

        matplotlib.pyplot.subplot(222)
        matplotlib.pyplot.scatter(i, torch.detach(model.lambda1), s=5, c="tab:blue")
        matplotlib.pyplot.subplot(224)
        matplotlib.pyplot.scatter(i, torch.detach(model.lambda2), s=5, c="tab:orange")

    matplotlib.pyplot.subplot(222)
    matplotlib.pyplot.xlabel(r"Epoch")
    matplotlib.pyplot.ylabel(r"$\lambda_1$ Estimated")

    matplotlib.pyplot.subplot(224)
    matplotlib.pyplot.xlabel(r"Epoch")
    matplotlib.pyplot.ylabel(r"$\lambda_2$ Estimated")

    matplotlib.pyplot.subplots_adjust(wspace=0.357, hspace=0.502)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
